// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>
#include <pollux/testing/dwio/batch_maker.h>
#include <pollux/common/file/temp_file_path.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/signature_binder.h>
#include <pollux/testing/expression/expression_fuzzer.h>
#include <pollux/testing/expression/fuzzer_runner.h>
#include <tests/expression/expression_runner.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/vector/vector_saver.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::test {

class ExpressionRunnerUnitTest : public testing::Test, public VectorBuilder {
 public:
  void SetUp() override {
    pollux::functions::prestosql::registerAllScalarFunctions();
  }

 protected:
  static void SetUpTestCase() {
    memory::MemoryManager::testingSetInstance({});
  }

  std::shared_ptr<memory::MemoryPool> pool_{
      memory::memoryManager()->addLeafPool()};
  std::shared_ptr<core::QueryCtx> queryCtx_{core::QueryCtx::create()};
  core::ExecCtx execCtx_{pool_.get(), queryCtx_.get()};
};

TEST_F(ExpressionRunnerUnitTest, run) {
  auto inputFile = exec::test::TempFilePath::create();
  auto selectivityVectorFile = exec::test::TempFilePath::create();
  auto resultFile = exec::test::TempFilePath::create();
  const auto inputPathStr = inputFile->getPath();
  const auto resultPathStr = resultFile->getPath();
  const auto selectivityVectorPathStr = selectivityVectorFile->getPath();
  const int vectorSize = 100;

  VectorMaker vectorMaker(pool_.get());
  auto inputVector = vectorMaker.row_vector(
      {"c0"}, {vectorMaker.flat_vector<StringView>(vectorSize, [](auto) {
        return "abc";
      })});
  SelectivityVector rows(vectorSize);
  auto resultVector = vectorMaker.row_vector(
      {"output0"}, {vectorMaker.flat_vector<int64_t>(vectorSize, [](auto row) {
        return row * 100;
      })});
  saveVectorToFile(inputVector.get(), inputPathStr.data());
  saveVectorToFile(resultVector.get(), resultPathStr.data());
  saveSelectivityVectorToFile(rows, selectivityVectorPathStr.data());

  for (bool useSeperatePoolForInput : {true, false}) {
    KLOG(INFO) << "Using useSeperatePoolForInput: " << useSeperatePoolForInput;
    EXPECT_NO_THROW(ExpressionRunner::run(
        inputPathStr.data(),
        selectivityVectorPathStr.data(),
        "length(c0)",
        "",
        resultPathStr.data(),
        "verify",
        0,
        "",
        "",
        nullptr,
        false,
        useSeperatePoolForInput));
  }
}

TEST_F(ExpressionRunnerUnitTest, persistAndReproComplexSql) {
  // Create a constant vector of ARRAY(Dictionary-Encoded INT)
  auto dictionary_vector = wrap_in_dictionary(
      make_indices({{2, 4, 0, 1}}), make_flat_vector<int32_t>({{1, 2, 3, 4, 5}}));
  auto arrVector = make_array_vector({0}, dictionary_vector, {});
  auto constantExpr = std::make_shared<core::ConstantTypedExpr>(
      BaseVector::wrap_in_constant(1, 0, arrVector));

  ASSERT_EQ(
      constantExpr->toString(),
      "4 elements starting at 0 {[0->2] 3, [1->4] 5, [2->0] 1, [3->1] 2}");

  auto sqlExpr = exec::ExprSet({constantExpr}, &execCtx_, false).expr(0);

  // Self contained SQL that flattens complex constant.
  auto selfContainedSql = sqlExpr->toSql();
  ASSERT_EQ(
      selfContainedSql,
      "ARRAY['3'::INTEGER, '5'::INTEGER, '1'::INTEGER, '2'::INTEGER]");

  std::vector<VectorPtr> complexConstants;
  auto complexConstantsSql = sqlExpr->toSql(&complexConstants);
  ASSERT_EQ(complexConstantsSql, "__complex_constant(c0)");

  auto row_vector = make_row_vector(complexConstants);

  // Emulate a reproduce from complex constant SQL
  auto sqlFile = exec::test::TempFilePath::create();
  auto complexConstantsFile = exec::test::TempFilePath::create();
  const auto complexConstantsFilePathStr = complexConstantsFile->getPath();
  auto sqlPathStr = sqlFile->getPath();
  auto sqlPath = sqlPathStr.c_str();
  auto complexConstantsPath = complexConstantsFilePathStr.c_str();

  // Write to file..
  saveStringToFile(complexConstantsSql, sqlPath);
  saveVectorToFile(row_vector.get(), complexConstantsPath);

  // Reproduce from file.
  auto reproSql = restoreStringFromFile(sqlPath);
  auto reproComplexConstants =
      restoreVectorFromFile(complexConstantsPath, pool_.get());

  auto reproExprs = ExpressionRunner::parseSql(
      reproSql, nullptr, pool_.get(), reproComplexConstants);
  ASSERT_EQ(reproExprs.size(), 1);
  // Note that ConstantExpr makes a copy of sharedConstantValue_ to guard
  // against race conditions, which in effect falttens the array.
  ASSERT_EQ(reproExprs[0]->toString(), "4 elements starting at 0 {3, 5, 1, 2}");
}

TEST_F(ExpressionRunnerUnitTest, primitiveConstantsInexpressibleInSql) {
  auto varbinaryData =
      vectorMaker_.flat_vector<StringView>({"12"_sv}, VARBINARY());
  auto constantExpr = std::make_shared<const core::ConstantTypedExpr>(
      BaseVector::wrap_in_constant(1, 0, varbinaryData));

  auto sqlExpr = exec::ExprSet({constantExpr}, &execCtx_, false).expr(0);

  ASSERT_THROW(sqlExpr->toSql(), PolluxUserError);
  std::vector<VectorPtr> complexConstants;
  auto complexConstantsSql = sqlExpr->toSql(&complexConstants);
  ASSERT_EQ(complexConstantsSql, "__complex_constant(c0)");
}
} // namespace kumo::pollux::test
